import torch
import grasp_cuda

def class_agnostic_nms(grasp_scores, grasp_preds, nms_config, score_thresh=None):
    src_grasp_scores = grasp_scores
    if score_thresh is not None:
        scores_mask = (grasp_scores >= score_thresh)
        grasp_scores = grasp_scores[scores_mask]
        grasp_preds = grasp_preds[scores_mask]

    selected = []
    if grasp_scores.shape[0] > 0:
        grasp_scores_nms, indices = torch.topk(grasp_scores, k=min(nms_config.NMS_PRE_MAXSIZE, grasp_scores.shape[0]))
        grasps_for_nms = grasp_preds[indices]
        #import time
        #torch.cuda. synchronize()
        #time_start = time.time()
        keep_idx, selected_scores = grasp_nms(
                grasps_for_nms[:, 0:10], grasp_scores_nms, nms_config.NMS_THRESH)
        #torch.cuda. synchronize()
        #print('time cost is : ', time.time() - time_start)
        #print(len(keep_idx))
        selected = indices[keep_idx[:nms_config.NMS_POST_MAXSIZE]]

#    if score_thresh is not None:
#        original_idxs = scores_mask.nonzero().view(-1)
#        selected = original_idxs[selected]
    return selected, src_grasp_scores[selected]

def grasp_nms(grasps, scores, thresh, pre_maxsize=None, **kwargs):
    assert grasps.shape[1] == 10
    order = scores.sort(0, descending=True)[1]
    if pre_maxsize is not None:
        order = order[:pre_maxsize]

    grasps = grasps[order].contiguous()
    keep = torch.LongTensor(grasps.size(0))
    num_out = grasp_cuda.grasp_nms_gpu(grasps, keep, thresh)
    return order[keep[:num_out].cuda()].contiguous(), None

def grasps_distance_gpu(grasps_a, grasps_b):
    assert grasps_a.shape[1] == grasps_b.shape[1] == 10


    distances_bev = torch.cuda.FloatTensor(torch.Size((grasps_a.shape[0], grasps_b.shape[0]))).zero_()  # (N, M)
    grasp_cuda.grasp_distance_gpu(grasps_a.contiguous(), grasps_b.contiguous(), distances_bev)


    return distances_bev
